"""Script demonstrating wrap_model_call middleware showing agent name for main and subagent calls."""

from langchain.agents import create_agent
from langchain.agents.middleware import wrap_model_call
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langgraph.config import get_config

# This will print the agent name for every model call (main or subagent)
@wrap_model_call
def print_agent_name(request, handler):
    """Middleware that prints the runtime agent_name and passes through the request."""
    print(f"[middleware] agent_name: {request.runtime.agent_name}")
    return handler(request)

# Define a subagent that will be called as a tool
subagent = create_agent(
    model="openai:gpt-4o-mini",
    middleware=[print_agent_name],
    name="subagent_graph"
)

# Define a tool that invokes the subagent
@tool
def trigger_subagent(query: str) -> str:
    """Call the subagent with the query."""
    response = subagent.invoke({"messages": [HumanMessage(content=query)]})
    # Return only the assistant's most recent message content
    return response["messages"][-1].content

# Main agent, which has access to the above tool
agent = create_agent(
    model="openai:gpt-4o-mini",
    tools=[trigger_subagent],
    middleware=[print_agent_name],
    name="sydney_graph"
)

# Compose a user message that requires the agent to call the subagent tool
result = agent.invoke({"messages": [HumanMessage(content="Use the trigger_subagent tool to answer: What's 2+2?")]})

print(result["messages"][-1].content)

# Example output:
# [middleware] agent_name: sydney_graph
# [middleware] agent_name: subagent_graph
# [middleware] agent_name: sydney_graph
# 2 + 2 equals 4.